# trace generated using paraview version 5.10.0
#
# import the simple module from the paraview
from paraview.simple import *
import os
from DEAD.AutoDecoder.ParaviewMacro.NameInterpreter import get_n_slots
import concurrent.futures

def show_one_BS(working_dir, id):
    n_slots = get_n_slots(id)
    theta = 360.0 / n_slots

    # create a new 'XML Structured Grid Reader'
    target_vts = XMLStructuredGridReader(
        FileName=[working_dir+"/"+id+".vts"])

    # get active view
    renderView1 = GetActiveViewOrCreate('RenderView')

    # show data in view
    target_vtsDisplay = Show(target_vts, renderView1)

    # trace defaults for the display properties.
    target_vtsDisplay.Representation = 'Outline'

    # reset view to fit data
    renderView1.ResetCamera()

    # update the view to ensure updated data information
    renderView1.Update()

    # hide data in view
    Hide(target_vts, renderView1)

    # create a new 'Iso Volume'
    isoVolume1 = IsoVolume(Input=target_vts)

    # Properties modified on isoVolume1
    isoVolume1.InputScalars = ['POINTS', 'define_marker_rho']
    isoVolume1.ThresholdRange = [0.0, 50.0]

    # show data in view
    isoVolume1Display = Show(isoVolume1, renderView1)

    # trace defaults for the display properties.
    isoVolume1Display.Representation = 'Surface'

    # hide data in view
    Hide(isoVolume1, renderView1)

    # create a new 'Iso Volume'
    isoVolume2 = IsoVolume(Input=isoVolume1)

    # Properties modified on isoVolume2
    isoVolume2.InputScalars = ['POINTS', 'define_marker_theta']
    isoVolume2.ThresholdRange = [0.0, 50.0]

    # show data in view
    isoVolume2Display = Show(isoVolume2, renderView1)

    # trace defaults for the display properties.
    isoVolume2Display.Representation = 'Surface'

    # hide data in view
    Hide(isoVolume2, renderView1)

    # create a new 'Contour'
    contour1 = Contour(Input=isoVolume2)

    # Properties modified on contour1
    contour1.Isosurfaces = [0.0]
    contour1.ComputeNormals = 1
    contour1.GenerateTriangles = 0

    # show data in view
    contour1Display = Show(contour1, renderView1)

    # trace defaults for the display properties.
    contour1Display.Representation = 'Surface'

    # turn off scalar coloring
    ColorBy(contour1Display, None)

    # update the view to ensure updated data information
    renderView1.Update()

    # create a new 'Reflect'
    reflect1 = Reflect(Input=contour1)

    # Properties modified on reflect1
    reflect1.Plane = 'Y Min'

    # show data in view
    reflect1Display = Show(reflect1, renderView1)

    # trace defaults for the display properties.
    reflect1Display.Representation = 'Surface'

    # turn off scalar coloring
    ColorBy(reflect1Display, None)

    # hide data in view
    Hide(contour1, renderView1)

    # update the view to ensure updated data information
    renderView1.Update()

    # change solid color
    reflect1Display.DiffuseColor = [1.0, 0.0, 0.0]

    transforms = []

    for i in range(n_slots-1):
        # create a new 'Transform'
        if i == 0:
            transform1 = Transform(Input=reflect1)
        else:
            transform1 = Transform(Input=transforms[i-1])

        # Properties modified on transform1.Transform
        transform1.Transform.Rotate = [0.0, 0.0, theta]

        transforms.append(transform1)

    for i in range(n_slots-1):
        # show data in view
        transform1Display = Show(transforms[i], renderView1)

        # trace defaults for the display properties.
        transform1Display.Representation = 'Surface'

        # turn off scalar coloring
        ColorBy(transform1Display, None)

        transform1Display.DiffuseColor = [1.0, 0.0, 0.0]

    # update the view to ensure updated data information
    renderView1.Update()

    # reset view to fit data
    renderView1.ResetCamera()

    # current camera placement for renderView1
    renderView1.CameraPosition = [-2, -2, -2]
    renderView1.CameraViewUp = [0.0, 1.0, 0.0]

    renderView1.ResetCamera()

    # Properties modified on renderView1
    renderView1.OrientationAxesVisibility = 0

    # save screenshot
    SaveScreenshot(working_dir+"/"+id+"BS.png", renderView1, ImageResolution=[2000, 2000],
                    TransparentBackground=1)
    print("finish BS visualization for "+id)

    # destroy
    for i in range(n_slots-1):
        Delete(transforms[i])
    del transforms

    Delete(reflect1)
    del reflect1

    Delete(contour1)
    del contour1

    Delete(isoVolume2)
    del isoVolume2

    Delete(isoVolume1)
    del isoVolume1

    Delete(target_vts)
    del target_vts


def show_all_BS(working_dir):
    # Get all files in the directory
    files = os.listdir(working_dir)

    # Filter files with .vts extension and get only the file names
    vts_file_names = [os.path.splitext(file)[0]
                      for file in files if file.endswith('.vts')]

    # Create a ProcessPoolExecutor
    with concurrent.futures.ProcessPoolExecutor(max_workers=3) as executor:
        # Submit tasks to the executor
        futures = [executor.submit(show_one_BS, working_dir, id) for id in vts_file_names]
    concurrent.futures.wait(futures)    

if __name__ == '__main__':
    from DEAD.AutoDecoder.Config import result_save_path
    show_all_BS(f"{result_save_path}/VTK")